'''
1.提示链（Prompt Chaining） 定义：将一个复杂任务分解为多个步骤，每个步骤使用一个提示（Prompt）引导模型生成中间结果，前一步的输出作为下一步的输入。

应用场景：
复杂推理任务：如数学解题、逻辑推理，分步骤引导模型“思考”。
内容生成流程：先生成大纲 → 再填充段落 → 最后润色。
多轮对话系统：用户输入 → 意图识别 → 信息提取 → 回复生成。
数据提取与结构化：从非结构化文本中逐步提取实体、关系、分类。
✅ 优势：提升模型在复杂任务上的准确性和可控性。
'''

from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from utils import init_llm, set_env, visualize_graph

# 初始化图，预定义存储的字段内容
class State(TypedDict):
    topic: str
    joke: str
    improved_joke: str
    final_joke: str


# 添加 LLM 调用节点
def generate_joke(state: State):
    """第一个 LLM 调用生成初始笑话"""

    msg = llm.invoke(f"写一个关于 {state['topic']} 话题的笑话/no_think")
    return {"joke": msg.content}


def improve_joke(state: State):
    """第二次调用 LLM 改进笑话"""

    msg = llm.invoke(f"通过添加文字游戏让这个笑话更有趣：{state['joke']}/no_think")
    return {"improved_joke": msg.content}


def polish_joke(state: State):
    """第三次 LLM 征集最终润色"""

    msg = llm.invoke(f"为这个笑话添加一个令人惊讶的转折： {state['improved_joke']}/no_think")
    return {"final_joke": msg.content}


def check_punchline(state: State):
    """检查笑话是否有妙语的门功能"""

    # 简单检查 - 笑话是否包含“？” 或 “！”
    if "?" in state["joke"] or "!" in state["joke"]:
        return "Pass"
    return "Fail"


def build_graph():
    workflow = StateGraph(State)

    workflow.add_node("generate_joke", generate_joke)
    workflow.add_node("improve_joke", improve_joke)
    workflow.add_node("polish_joke", polish_joke)

    workflow.add_edge(START, "generate_joke")
    workflow.add_conditional_edges(
        "generate_joke", check_punchline, {"Fail": "improve_joke", "Pass": END}
    )
    workflow.add_edge("improve_joke", "polish_joke")
    workflow.add_edge("polish_joke", END)

    chain = workflow.compile()

    return chain


if __name__ == "__main__":
    set_env()  # 设置环境变量
    llm = init_llm(temperature=0)  # 初始化LLM，温度设为0以获得确定性输出
    chain = build_graph()  # 构建工作流图

    visualize_graph(chain)

    state = chain.invoke({"topic": "猫"})
    print("Initial joke:")
    print(state["joke"])
    print("\n--- --- ---\n")
    if "improved_joke" in state:
        print("Improved joke:")
        print(state["improved_joke"])
        print("\n--- --- ---\n")

        print("Final joke:")
        print(state["final_joke"])
    else:
        print("笑话质量门失败 - 未检测到妙语！")